from QKDNetwork import QKDNetwork
from noise_provision import provision_with_greedy
from lb_noise import calculateFlowPath, estimateSlotNum
import random
from copy import deepcopy

# QKD请求的路由方案数据结构  [{"sd": [1, 2], "path": [1, 3, 2]}...]


class Compare:
    def __init__(self, net: QKDNetwork, avg_noise, threshold_noise):
        self.net: QKDNetwork = net
        self.avg_noise = avg_noise  # 平均噪声 0-19
        self.threshold_noise = threshold_noise  # 噪声阈值 0-19
        self.noised_qkd_networks = []
        self.generated_noise_networks()

    # 生成500个噪声网络
    def generated_noise_networks(self):
        for i in range(500):
            net = deepcopy(self.net)
            self.addNoise(net)
            self.noised_qkd_networks.append(net)

    # 向网络中注入噪声
    def inject_noise(self, net: QKDNetwork, idx):
        for node in net.G.nodes:
            net.G.nodes[node]["transmitter_rate"] = self.noised_qkd_networks[
                idx
            ].G.nodes[node]["transmitter_rate"]

    def getData(self):
        # return [self.rpsp_noise_threshold()]  # 考虑噪声和噪声阈值
        return [
            self.rpsp(),
            self.rpsp_noise_threshold(),
            self.rpsp_noise(),
        ]  # 不考虑噪声和考虑噪声

    def is_terminate_schedule(self, net: QKDNetwork):
        flag = True
        for sd_index, sd_item in enumerate(net.sd_list):
            # 还有没满足的SD需求，继续调度
            if sd_item[2] > sd_item[3]:
                flag = False
                break
        return flag

    def addNoise(self, net):
        # 向链路增加噪声
        for node in net.G.nodes:
            random_number = random.random()  # 0 ~ 1

            net.G.nodes[node]["transmitter_rate"] *= (
                1 - self.avg_noise * 2.5 / 100
            ) * (1.05 - random_number / 10)
            if net.G.nodes[node]["transmitter_rate"] < 0:
                net.G.nodes[node]["transmitter_rate"] = 0.1

    def get_sd_idx_by_sd_pair(self, sd_start, sd_end):
        for i in range(len(self.net.sd_list)):
            if self.net.sd_list[i][0] == sd_start and self.net.sd_list[i][1] == sd_end:
                return i

    def rpsp(self, with_noise=False, with_threshold=False):
        terminate_schedule_flag = False
        max_time_slot = 0

        tmp_net: QKDNetwork = deepcopy(self.noised_qkd_networks[0])

        route = None
        p_normalized = None
        t_value = None  # 上次算出来的最优的时隙数量
        init_t_value = None  # 初次算出来的最优的时隙数量

        while not terminate_schedule_flag:
            max_time_slot += 1
            # print(f"\rmax_time_slot: {max_time_slot}", end="", flush=True)
            # 在一个时隙内，算路、算PS供应、在SD对上记账、在链路上记账

            if (route is None) or with_noise:
                if max_time_slot == 1:
                    init_t_value = estimateSlotNum(tmp_net)
                    # print("init time estimate", t_value)
                route, p_normalized, _ = calculateFlowPath(tmp_net)
                # print(f"t{max_time_slot}", t_value + max_time_slot - 1)

            if with_threshold and max_time_slot > 1:
                # 当前估计出来的总时隙值
                tmp_t_value = estimateSlotNum(tmp_net) + max_time_slot - 1
                # print("tmp_t_value", tmp_t_value)
                if tmp_t_value - init_t_value > self.threshold_noise:
                    route, p_normalized, _ = calculateFlowPath(tmp_net)

            # 取出来一些路，发射密钥（一个时隙内）
            provision_with_greedy(route, tmp_net, p_normalized)

            total_remain_key_num = 0
            for edge in tmp_net.G.edges:
                total_remain_key_num += tmp_net.G[edge[0]][edge[1]]["remain_key_num"]

            terminate_schedule_flag = self.is_terminate_schedule(tmp_net)
            self.inject_noise(tmp_net, max_time_slot)
            if max_time_slot > 200:
                break
        # print("final actual time", max_time_slot)
        # print(f"final time slot{' optimal' if with_noise else ''}:", max_time_slot)
        return max_time_slot

    def rpsp_noise(self):
        return self.rpsp(with_noise=True)

    def rpsp_noise_threshold(self):
        # 考虑噪声阈值
        return self.rpsp(with_noise=False, with_threshold=True)
